import os
import sys
import time
import tempfile
import hashlib
import threading
import re
import numpy as np
from PIL import Image
import pyperclip
import pytesseract
import keyboard
from PyQt5.QtWidgets import QApplication, QSystemTrayIcon, QMenu, QAction, QMessageBox
from PyQt5.QtGui import QIcon, QImage
from PyQt5.QtCore import QObject, pyqtSignal, QTimer

# 配置Tesseract路径（根据你的安装路径修改）
pytesseract.pytesseract.tesseract_cmd = r'D:\Program Files\Tesseract-OCR\tesseract.exe'

# 获取当前脚本所在目录（推荐）
project_dir = os.path.dirname(os.path.abspath(__file__))  # 当前脚本所在文件夹
# 定义子文件夹用于存放剪贴板图片
clipboard_folder = os.path.join(project_dir, "clipboard_images")
# 确保目录存在
os.makedirs(clipboard_folder, exist_ok=True)

class ClipboardMonitor(QObject):
    update_signal = pyqtSignal(str)

    def __init__(self):
        super().__init__()
        self.last_hash = None
        self.temp_dir = tempfile.gettempdir()
        self.running = True
        self.paused = False

        # 敏感信息正则表达式模式
        self.sensitive_patterns = [
            # 密码模式（8位以上包含数字和字母）
            re.compile(r'(?<![a-zA-Z0-9])([a-zA-Z0-9!@#$%^&*()_+=\-\[\]{}|;:,.<>?/~`]{8,})(?![a-zA-Z0-9])'),
            # 邮箱模式
            re.compile(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b'),
            # 手机号模式
            re.compile(r'(?<!\d)(1[3-9]\d{9})(?!\d)'),
            # 身份证号模式
            re.compile(r'[1-9]\d{5}(18|19|20)\d{2}(0[1-9]|1[0-2])(0[1-9]|[12]\d|3[01])\d{3}[\dXx]'),
            # 银行卡号模式
            re.compile(r'(?<!\d)(\d{16,19})(?!\d)')
        ]

        # 设置系统托盘
        self.setup_tray()

        # 启动监控线程
        self.monitor_thread = threading.Thread(target=self.monitor_clipboard, daemon=True)
        self.monitor_thread.start()

        # 设置定时器更新状态
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_status)
        self.timer.start(1000)  # 每秒更新一次状态

    def setup_tray(self):
        self.app = QApplication(sys.argv)
        self.tray_icon = QSystemTrayIcon(QIcon())
        self.tray_icon.setToolTip("剪贴板敏感信息脱敏工具")

        # 创建托盘菜单
        tray_menu = QMenu()

        self.pause_action = QAction("暂停监控", self)
        self.pause_action.triggered.connect(self.toggle_pause)
        tray_menu.addAction(self.pause_action)

        exit_action = QAction("退出", self)
        exit_action.triggered.connect(self.exit_app)
        tray_menu.addAction(exit_action)

        self.tray_icon.setContextMenu(tray_menu)
        self.tray_icon.show()

        # 显示通知
        self.tray_icon.showMessage(
            "剪贴板敏感信息脱敏工具",
            "工具已启动，正在监控剪贴板...",
            QSystemTrayIcon.Information,
            3000
        )

    def toggle_pause(self):
        self.paused = not self.paused
        if self.paused:
            self.pause_action.setText("继续监控")
            self.tray_icon.showMessage(
                "监控已暂停",
                "剪贴板监控已暂停，点击继续监控恢复",
                QSystemTrayIcon.Information,
                2000
            )
        else:
            self.pause_action.setText("暂停监控")
            self.tray_icon.showMessage(
                "监控已恢复",
                "剪贴板监控已恢复",
                QSystemTrayIcon.Information,
                2000
            )

    def update_status(self):
        if self.paused:
            self.tray_icon.setToolTip("剪贴板敏感信息脱敏工具 (已暂停)")
        else:
            self.tray_icon.setToolTip("剪贴板敏感信息脱敏工具 (运行中)")

    def exit_app(self):
        self.running = False
        self.tray_icon.hide()
        QApplication.quit()
        sys.exit(0)

    def monitor_clipboard(self):
        while self.running:
            if not self.paused:
                try:
                    # 获取剪贴板内容
                    content = pyperclip.paste()

                    # 检查是否为图片
                    if isinstance(content, Image.Image):
                        # 计算图片哈希，避免重复处理
                        current_hash = hashlib.md5(content.tobytes()).hexdigest()

                        if current_hash != self.last_hash:
                            self.last_hash = current_hash
                            self.process_image(content)
                except Exception as e:
                    print(f"监控剪贴板时出错: {e}")

            time.sleep(0.5)  # 每0.5秒检查一次

    def process_image(self, image):
        try:
            # 保存图片到临时目录
            # 设置为 self.temp_dir
            self.temp_dir = clipboard_folder
            temp_path = os.path.join(self.temp_dir, f"clipboard_{int(time.time())}.png")
            image.save(temp_path)

            # OCR识别图片中的文字
            data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)

            # 查找敏感信息
            sensitive_regions = []
            for i in range(len(data['text'])):
                text = data['text'][i]
                if not text.strip():
                    continue

                # 检查是否匹配敏感信息模式
                for pattern in self.sensitive_patterns:
                    for match in pattern.finditer(text):
                        start, end = match.span()
                        # 获取敏感信息在图片中的位置
                        x = data['left'][i] + int(data['width'][i] * start / len(text))
                        y = data['top'][i]
                        width = int(data['width'][i] * (end - start) / len(text))
                        height = data['height'][i]

                        sensitive_regions.append((x, y, width, height))

            # 如果有敏感信息，进行马赛克处理
            if sensitive_regions:
                self.apply_mosaic(image, sensitive_regions)

                # 将处理后的图片放回剪贴板
                self.copy_image_to_clipboard(image)

                # 发送通知
                self.update_signal.emit(f"检测到敏感信息并已脱敏处理 ({len(sensitive_regions)}处)")
                self.tray_icon.showMessage(
                    "敏感信息已脱敏",
                    f"检测到 {len(sensitive_regions)} 处敏感信息并已进行马赛克处理",
                    QSystemTrayIcon.Warning,
                    3000
                )
        except Exception as e:
            print(f"处理图片时出错: {e}")

    def apply_mosaic(self, image, regions, cell_size=10):
        """对图片中的指定区域应用马赛克效果"""
        img_array = np.array(image)

        for x, y, w, h in regions:
            # 确保区域在图片范围内
            x = max(0, x)
            y = max(0, y)
            w = min(w, image.width - x)
            h = min(h, image.height - y)

            if w <= 0 or h <= 0:
                continue

            # 提取区域
            region = img_array[y:y + h, x:x + w]

            # 计算马赛克块大小
            if w < cell_size or h < cell_size:
                cell_size = min(w, h)

            # 计算马赛克块数量
            h_blocks = h // cell_size
            w_blocks = w // cell_size

            # 应用马赛克
            for i in range(h_blocks):
                for j in range(w_blocks):
                    # 获取当前块
                    block = region[i * cell_size:(i + 1) * cell_size, j * cell_size:(j + 1) * cell_size]
                    # 计算块的平均颜色
                    avg_color = np.mean(block, axis=(0, 1))
                    # 将整个块设置为平均颜色
                    region[i * cell_size:(i + 1) * cell_size, j * cell_size:(j + 1) * cell_size] = avg_color

            # 处理边缘部分
            if h % cell_size != 0:
                block = region[h_blocks * cell_size:h, :]
                avg_color = np.mean(block, axis=(0, 1))
                region[h_blocks * cell_size:h, :] = avg_color

            if w % cell_size != 0:
                block = region[:, w_blocks * cell_size:w]
                avg_color = np.mean(block, axis=(0, 1))
                region[:, w_blocks * cell_size:w] = avg_color

        # 将处理后的数组转换回图片
        processed_image = Image.fromarray(img_array)
        return processed_image

    def copy_image_to_clipboard(self, image):
        """将图片复制到剪贴板"""
        try:
            # 将PIL图片转换为QImage
            qimage = QImage(image.tobytes(), image.width, image.height, image.width * 3, QImage.Format_RGB888)

            # 复制到剪贴板
            clipboard = QApplication.clipboard()
            clipboard.setImage(qimage)
        except Exception as e:
            print(f"复制图片到剪贴板时出错: {e}")


def main():
    # 创建并启动监控器
    monitor = ClipboardMonitor()

    # 设置全局热键 Ctrl+Alt+D 退出程序
    # keyboard.add_hotkey('ctrl+alt+d', monitor.exit_app)

    # 运行Qt应用
    sys.exit(monitor.app.exec_())


if __name__ == "__main__":
    main()